# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# A script to run multinode training with submitit.
# --------------------------------------------------------

import argparse
import os
import uuid
from pathlib import Path

import main_pretrain as trainer
import submitit


def parse_args():
    trainer_parser = trainer.get_args_parser()
    parser = argparse.ArgumentParser("Submitit for MAE pretrain", parents=[trainer_parser])
    parser.add_argument("--ngpus", default=8, type=int, help="Number of gpus to request on each node")
    parser.add_argument("--nodes", default=2, type=int, help="Number of nodes to request")
    parser.add_argument("--timeout", default=4320, type=int, help="Duration of the job")
    parser.add_argument("--job_dir", default="", type=str, help="Job dir. Leave empty for automatic.")

    parser.add_argument("--partition", default="learnfair", type=str, help="Partition where to submit")
    parser.add_argument("--use_volta32", action='store_true', help="Request 32G V100 GPUs")
    parser.add_argument('--comment', default="", type=str, help="Comment to pass to scheduler")
    return parser.parse_args()


def get_shared_folder() -> Path:
    """
    获取一个共享文件夹的路径。如果 /checkpoint 目录存在，则创建并返回
    当前用户下的实验文件夹路径。如果 /checkpoint 目录不存在，抛出异常。
    """
    user = os.getenv("USER")  # 获取当前操作系统中的用户名
    if Path("/checkpoint/").is_dir():  # 检查 /checkpoint 目录是否存在
        p = Path(f"/checkpoint/{user}/experiments")  # 获取当前用户的实验文件夹路径
        p.mkdir(exist_ok=True)  # 如果该文件夹不存在，则创建它
        return p  # 返回文件夹路径
    raise RuntimeError("No shared folder available")  # 如果 /checkpoint 目录不存在，抛出异常


def get_init_file():
    """
    创建一个初始化文件。初始化文件的父目录必须存在，文件本身不能已经存在。
    每次调用时都会返回一个新的初始化文件路径，文件名包含 UUID 确保唯一性。
    """
    os.makedirs(str(get_shared_folder()), exist_ok=True)  # 确保共享文件夹已创建
    init_file = get_shared_folder() / f"{uuid.uuid4().hex}_init"  # 使用 UUID 创建唯一的初始化文件路径
    if init_file.exists():  # 如果文件已经存在
        os.remove(str(init_file))  # 删除旧的文件
    return init_file  # 返回新创建的初始化文件路径


class Trainer(object):
    """
    训练器类，负责训练过程的管理，包括 GPU 参数配置、模型训练执行等。
    """

    def __init__(self, args):
        """
        初始化 Trainer 实例，保存训练配置参数。

        参数:
            args: 包含训练参数和超参数的命令行参数或配置字典。
        """
        self.args = args  # 保存传入的训练参数

    def __call__(self):
        """
        调用训练器的训练流程，执行预训练任务。
        """
        import main_pretrain as trainer  # 导入预训练脚本

        self._setup_gpu_args()  # 配置与 GPU 相关的参数
        trainer.main(self.args)  # 调用预训练脚本的 main 函数开始训练

    def checkpoint(self):
        """
        检查点功能，如果存在检查点文件，则恢复训练。如果没有，则继续训练。
        返回一个延迟任务，用于重新排队训练任务。
        """
        import os
        import submitit

        # 获取初始化文件，并设置 dist_url 用于分布式训练
        self.args.dist_url = get_init_file().as_uri()
        checkpoint_file = os.path.join(self.args.output_dir, "checkpoint.pth")  # 设置检查点路径
        if os.path.exists(checkpoint_file):  # 如果检查点文件存在
            self.args.resume = checkpoint_file  # 设置恢复路径

        print("Requeuing ", self.args)  # 打印重新排队的参数
        empty_trainer = type(self)(self.args)  # 创建一个新的 Trainer 实例
        return submitit.helpers.DelayedSubmission(empty_trainer)  # 返回一个延迟提交任务

    def _setup_gpu_args(self):
        """
        配置与 GPU 和分布式训练相关的参数，包括作业环境和资源分配。
        """
        import submitit
        from pathlib import Path

        job_env = submitit.JobEnvironment()  # 获取作业环境信息
        # 设置输出目录，替换作业 ID 占位符为当前作业 ID
        self.args.output_dir = Path(str(self.args.output_dir).replace("%j", str(job_env.job_id)))
        self.args.log_dir = self.args.output_dir  # 将日志目录设置为输出目录
        self.args.gpu = job_env.local_rank  # 当前进程对应的 GPU 编号
        self.args.rank = job_env.global_rank  # 当前进程的全局编号（在分布式训练中使用）
        self.args.world_size = job_env.num_tasks  # 总的任务数量（作业并行度）
        print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")  # 打印任务组信息


def main():
    """
    主函数，解析命令行参数并提交训练任务。使用 `submitit` 调度作业并配置资源。
    """
    args = parse_args()  # 解析命令行参数
    if args.job_dir == "":  # 如果没有指定 job_dir，则使用默认共享文件夹路径
        args.job_dir = get_shared_folder() / "%j"  # 使用作业 ID 生成独立的目录

    # 使用作业 ID 来创建独立的文件夹以便追踪实验
    executor = submitit.AutoExecutor(folder=args.job_dir, slurm_max_num_timeout=30)  # 创建执行器实例

    # 配置资源参数
    num_gpus_per_node = args.ngpus  # 每个节点的 GPU 数量
    nodes = args.nodes  # 使用的节点数
    timeout_min = args.timeout  # 超时时间（分钟）

    partition = args.partition  # SLURM 分区
    kwargs = {}  # 额外的参数
    if args.use_volta32:  # 如果使用特定的硬件（如 Volta 32GB GPU），设置约束
        kwargs['slurm_constraint'] = 'volta32gb'
    if args.comment:  # 如果提供了作业注释，添加注释信息
        kwargs['slurm_comment'] = args.comment

    # 更新执行器参数，设置内存、GPU、CPU 等资源要求
    executor.update_parameters(
        mem_gb=40 * num_gpus_per_node,  # 每个节点所需的内存（GB）
        gpus_per_node=num_gpus_per_node,  # 每个节点的 GPU 数量
        tasks_per_node=num_gpus_per_node,  # 每个 GPU 分配一个任务
        cpus_per_task=10,  # 每个任务分配的 CPU 核心数
        nodes=nodes,  # 节点数量
        timeout_min=timeout_min,  # 超时时间（分钟）
        # 下面是集群特定的 SLURM 参数
        slurm_partition=partition,  # 使用的 SLURM 分区
        slurm_signal_delay_s=120,  # SLURM 信号延迟时间（秒）
        **kwargs  # 额外的参数
    )

    executor.update_parameters(name="mae")  # 设置作业名称为 "mae"

    # 获取初始化文件并设置分布式训练的相关参数
    args.dist_url = get_init_file().as_uri()
    args.output_dir = args.job_dir  # 输出目录

    # 创建 Trainer 实例并提交作业
    trainer = Trainer(args)
    job = executor.submit(trainer)  # 提交训练作业

    # 打印作业 ID
    print(job.job_id)


if __name__ == "__main__":
    main()  # 调用主函数，开始执行
